﻿using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CodeActions;
using Microsoft.CodeAnalysis.CodeRefactorings;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using Microsoft.CodeAnalysis.Formatting;
using Microsoft.CodeAnalysis.Simplification;
using Microsoft.CodeAnalysis.Text;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

namespace CodeAnalysisRuleCollection.Enumerables.CSharp
{
    /// <example>source.ElementAt(source.Count() - 1)</example>
    [ExportCodeRefactoringProvider("ElementAccess", LanguageNames.CSharp)]
    internal partial class ElementAccessRefactoringProvider : ICodeRefactoringProvider
    {
        public async Task<IEnumerable<CodeAction>> GetRefactoringsAsync(Document document, TextSpan span, CancellationToken cancellationToken)
        {
            var root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false) as CompilationUnitSyntax;
            var model = await document.GetSemanticModelAsync(cancellationToken).ConfigureAwait(false);

            IList<CodeAction> actions = new List<CodeAction>();

            // ElementAt()
            if ((var invocation = root.FindToken(span.Start).Parent.AncestorsAndSelf().OfType<InvocationExpressionSyntax>().FirstOrDefault()) != null)
            {
                if (IsCallToEnumerableElementAt(model, invocation, cancellationToken) &&
                    invocation.ArgumentList.Arguments.Count == 1
                )
                {
                    // ElementAt(0)
                    if (IsZeroConstantValue(model, invocation.ArgumentList.Arguments.Single().Expression, cancellationToken))
                        actions.Add(CodeAction.Create("Change use of Enumerable.ElementAt() to Enumerable.First()", c => ChangeElementAtToFirstAsync(root, document, invocation, c)));

                    MemberAccessExpressionSyntax memberAccess = invocation.Expression as MemberAccessExpressionSyntax;

                    // check if it has an indexer
                    var typeInfo = model.GetTypeInfo(memberAccess.Expression, cancellationToken);

                    if (typeInfo.Type != null &&
                        (typeInfo.Type.TypeKind == TypeKind.ArrayType ||
                        typeInfo.Type.Interfaces.Contains(model.Compilation.GetTypeByMetadataName(typeof(IList<>).FullName)))
                    )
                    {
                        actions.Add(CodeAction.Create("Change use of Enumerable.ElementAt() to element access", c => ChangeElementAtToIndexerAsync(root, document, invocation, c)));

                        // ElementAt(Count - 1)
                        if (IsLengthOrCountMinusOne(model, memberAccess.Expression, invocation.ArgumentList.Arguments.Single().Expression, cancellationToken))
                            actions.Add(CodeAction.Create("Change use of element access to Enumerable.Last()", c => ChangeElementAtToLastAsync(root, document, invocation, c)));
                    }
                }
            }

            cancellationToken.ThrowIfCancellationRequested();

            // array[]
            if ((var elementAccess = root.FindToken(span.Start).Parent.AncestorsAndSelf().OfType<ElementAccessExpressionSyntax>().FirstOrDefault()) != null)
            {
                if (elementAccess.ArgumentList.Arguments.Count == 1)
                {
                    // array[0]
                    if (IsZeroConstantValue(model, elementAccess.ArgumentList.Arguments.Single().Expression, cancellationToken))
                        actions.Add(CodeAction.Create("Change use of element access to Enumerable.First()", c => ChangeIndexerToFirstAsync(root, document, elementAccess, c)));

                    // array[expr]
                    actions.Add(CodeAction.Create("Change use of element access to Enumerable.ElementAt()", c => ChangeIndexerToElementAtAsync(root, document, elementAccess, c)));

                    // array[array.Length - 1]
                    if (IsLengthOrCountMinusOne(model, elementAccess.Expression, elementAccess.ArgumentList.Arguments.Single().Expression, cancellationToken))
                        actions.Add(CodeAction.Create("Change use of element access to Enumerable.Last()", c => ChangeIndexerToLastAsync(root, document, elementAccess, c)));
                }
            }

            return actions.Any() ? actions : null;
        }

        private bool IsCallToEnumerableElementAt(SemanticModel semanticModel, ExpressionSyntax expression, CancellationToken cancellationToken)
        {
            var invocation = expression as InvocationExpressionSyntax;
            if (invocation == null)
                return false;

            var call = invocation.Expression as MemberAccessExpressionSyntax;
            if (call == null)
                return false;

            var methodSymbol = semanticModel.GetSymbolInfo(call, cancellationToken).Symbol as IMethodSymbol;
            if (methodSymbol == null ||
                methodSymbol.Name != "ElementAt" ||
                methodSymbol.ConstructedFrom == null)
                return false;

            var enumerable = semanticModel.Compilation.GetTypeByMetadataName(typeof(Enumerable).FullName);
            if (enumerable == null || !methodSymbol.ConstructedFrom.ContainingType.Equals(enumerable))
                return false;

            return true;
        }

        private bool IsZeroConstantValue(SemanticModel semanticModel, ExpressionSyntax expression, CancellationToken cancellationToken)
        {
            var constantValue = semanticModel.GetConstantValue(expression, cancellationToken);

            return constantValue.HasValue && constantValue.Value.Equals(0);
        }

        private bool IsLengthOrCountMinusOne(SemanticModel semanticModel, ExpressionSyntax member, ExpressionSyntax expression, CancellationToken cancellationToken)
        {
            // a - 1
            BinaryExpressionSyntax binary = expression as BinaryExpressionSyntax;
            if (binary == null)
                return false;

            if (!binary.OperatorToken.IsKind(SyntaxKind.MinusToken))
                return false;

            // 1
            var constant = semanticModel.GetConstantValue(binary.Right, cancellationToken);
            if (!constant.HasValue || !constant.Value.Equals(1))
                return false;

            // source.Length, source.Count
            if ((var memberAccess = binary.Left as MemberAccessExpressionSyntax) != null)
            {
                if (!memberAccess.Expression.IsEquivalentTo(member))
                    return false;

                IPropertySymbol propertySymbol = semanticModel.GetSymbolInfo(memberAccess, cancellationToken).Symbol as IPropertySymbol;

                if (propertySymbol == null)
                    return false;

                if ((propertySymbol.Name == "Length" || propertySymbol.Name == "LongLength") && propertySymbol.ContainingType.Equals(semanticModel.Compilation.GetTypeByMetadataName(typeof(Array).FullName)) ||
                    propertySymbol.Name == "Count" && propertySymbol.ContainingType.AllInterfaces.Contains(semanticModel.Compilation.GetTypeByMetadataName(typeof(System.Collections.ICollection).FullName)))
                    return true;

                return false;
            }

            // source.Count()
            if ((var invocation = binary.Left as InvocationExpressionSyntax) != null)
            {
                var countMember = invocation.Expression as MemberAccessExpressionSyntax;
                if (countMember == null)
                    return false;

                if (!countMember.Expression.IsEquivalentTo(member))
                    return false;

                IMethodSymbol methodSymbol = semanticModel.GetSymbolInfo(countMember, cancellationToken).Symbol as IMethodSymbol;

                if (methodSymbol == null)
                    return false;

                if (methodSymbol.Name == "Count")
                {
                    // not Count(...)
                    if (invocation.ArgumentList.Arguments.Count != 0)
                        return false;

                    if (methodSymbol.ConstructedFrom == null)
                        return false;

                    var enumerable = semanticModel.Compilation.GetTypeByMetadataName(typeof(Enumerable).FullName);

                    if (enumerable == null ||
                        !methodSymbol.ConstructedFrom.ContainingType.Equals(enumerable))
                        return false;

                    return true;
                }
                else if (methodSymbol.Name == "GetLength")
                {
                    IArrayTypeSymbol arraySymbol = semanticModel.GetTypeInfo(member, cancellationToken).Type as IArrayTypeSymbol;

                    if (arraySymbol == null || arraySymbol.Rank != 1)
                        return false;

                    // should be GetLength(0)
                    if (invocation.ArgumentList.Arguments.Count != 1)
                        return false;

                    Optional<object> constantValue = semanticModel.GetConstantValue(invocation.ArgumentList.Arguments.Single().Expression, cancellationToken);

                    if (!constantValue.HasValue || !constantValue.Value.Equals(0))
                        return false;

                    if (methodSymbol.ConstructedFrom == null)
                        return false;

                    if (!methodSymbol.ConstructedFrom.ContainingType.Equals(semanticModel.Compilation.GetTypeByMetadataName(typeof(Array).FullName)))
                        return false;

                    return true;
                }

                return false;
            }

            return false;
        }

        private async Task<Document> ChangeElementAtToFirstAsync(SyntaxNode root, Document document, InvocationExpressionSyntax expression, CancellationToken cancellationToken)
        {
            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(expression,
                expression
                    .WithExpression(
                        (expression.Expression as MemberAccessExpressionSyntax).WithName(SyntaxFactory.IdentifierName("First"))
                    )
                    .WithArgumentList(SyntaxFactory.ArgumentList())
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeElementAtToLastAsync(SyntaxNode root, Document document, InvocationExpressionSyntax expression, CancellationToken cancellationToken)
        {
            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(expression,
                expression
                    .WithExpression(
                        (expression.Expression as MemberAccessExpressionSyntax).WithName(SyntaxFactory.IdentifierName("Last"))
                    )
                    .WithArgumentList(SyntaxFactory.ArgumentList())
            );

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeIndexerToElementAtAsync(SyntaxNode root, Document document, ElementAccessExpressionSyntax expression, CancellationToken cancellationToken)
        {
            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(expression,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, expression.Expression, SyntaxFactory.IdentifierName("ElementAt")),
                    SyntaxFactory.ArgumentList(expression.ArgumentList.Arguments)
                )
                    .WithAdditionalAnnotations(Simplifier.Annotation)
            );

            newRoot = (newRoot as CompilationUnitSyntax).WithUsing("System.Linq");

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeIndexerToFirstAsync(SyntaxNode root, Document document, ElementAccessExpressionSyntax expression, CancellationToken cancellationToken)
        {
            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(expression,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, expression.Expression, SyntaxFactory.IdentifierName("First"))
                )
                    .WithAdditionalAnnotations(Simplifier.Annotation)
            );

            newRoot = (newRoot as CompilationUnitSyntax).WithUsing("System.Linq");

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeIndexerToLastAsync(SyntaxNode root, Document document, ElementAccessExpressionSyntax expression, CancellationToken cancellationToken)
        {
            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(expression,
                SyntaxFactory.InvocationExpression(
                    SyntaxFactory.MemberAccessExpression(SyntaxKind.SimpleMemberAccessExpression, expression.Expression, SyntaxFactory.IdentifierName("Last"))
                )
                    .WithAdditionalAnnotations(Simplifier.Annotation)
            );

            newRoot = (newRoot as CompilationUnitSyntax).WithUsing("System.Linq");

            return document.WithSyntaxRoot(newRoot);
        }

        private async Task<Document> ChangeElementAtToIndexerAsync(SyntaxNode root, Document document, InvocationExpressionSyntax expression, CancellationToken cancellationToken)
        {
            var expr = (expression.Expression as MemberAccessExpressionSyntax).Expression;

            // take care of new int[0] -> new int[0][idx]
            if (expr.IsKind(SyntaxKind.ArrayCreationExpression))
                expr = SyntaxFactory.ParenthesizedExpression(expr); // (new int[0])[expr]

            // replace node
            var newRoot = root.ReplaceNode<SyntaxNode, ExpressionSyntax>(expression,
                SyntaxFactory.ElementAccessExpression(
                    expr, SyntaxFactory.BracketedArgumentList(expression.ArgumentList.Arguments)
                )
            );

            return document.WithSyntaxRoot(newRoot);
        }
    }

    public static class CompilationUnitSyntaxExtensions
    {
        public static CompilationUnitSyntax WithUsing(this CompilationUnitSyntax root, string name)
        {
            if (!root.Usings.Any(u => u.Name.ToString() == name))
                return root.AddUsings(SyntaxFactory.UsingDirective(SyntaxFactory.ParseName(name))
                    .WithAdditionalAnnotations(Formatter.Annotation));

            return root;
        }
    }
}
